#include <iostream>
#include <deque>
#include <opencv2/core.hpp>
#include <opencv2/imgcodecs.hpp>
#include <opencv2/imgproc.hpp>
#include <opencv2/highgui.hpp>

// step 1
cv::Mat _gaussian_filter(const cv::Mat& mat) {
    cv::Mat matDouble;
    mat.convertTo(matDouble, CV_64FC1);
    cv::Mat kernel = (cv::Mat_<double>(5, 5) << 
        2, 4, 5, 4, 2,
        4, 9, 12, 9, 4,
        5, 12, 15, 12, 5,
        4, 9, 12, 9, 4,
        2, 4, 5, 4, 2);
    kernel = kernel/159;
    cv::Mat resDouble;
    cv::filter2D(matDouble, resDouble, -1, kernel, cv::Point(-1, -1), 0.0, cv::BORDER_REFLECT101);
    cv::Mat res;
    resDouble.convertTo(res, CV_8UC1);
    return res;
}

// step 2
void _sobel_gradient(const cv::Mat& mat, cv::Mat& dx, cv::Mat& dy, cv::Mat& magnitudes, cv::Mat& angles, int apertureSize, bool L2gradient) {
    CV_Assert(apertureSize == 3 || apertureSize == 5);
    
    double scale = 1.0;
    cv::Sobel(mat, dx, CV_16S, 1, 0, apertureSize, scale, 0, cv::BORDER_REPLICATE);
    cv::Sobel(mat, dy, CV_16S, 0, 1, apertureSize, scale, 0, cv::BORDER_REPLICATE);
    
    const int TAN225 = 13573;
    
    angles = cv::Mat(mat.size(), CV_8UC1);  // 0-> horizontal, 1 -> vertical, 2 -> diagonal
    magnitudes = cv::Mat::zeros(mat.rows+2, mat.cols+2, CV_32SC1);
    cv::Mat magROI = cv::Mat(magnitudes, cv::Rect(1, 1, mat.cols, mat.rows));
    for (int i = 0; i < mat.rows; i++) {
        for (int j = 0; j < mat.cols; j++) {
            short xs = dx.ptr<short>(i)[j];
            short ys = dy.ptr<short>(i)[j];
            int x = (int)std::abs(xs);
            int y = (int)std::abs(ys) << 15;
            
            if (L2gradient) {
                magROI.ptr<int>(i)[j] = int(xs) * xs + int(ys) * ys;
            } else {
                magROI.ptr<int>(i)[j] = std::abs(int(xs)) + std::abs(int(ys));
            }
            
            int tan225x = x * TAN225;
            if (y < tan225x) {  // horizontal
                angles.ptr<uchar>(i)[j] = 0;
            } else {
                int tan675x = tan225x + (x << 16);
                if (y > tan675x) {  // vertical
                    angles.ptr<uchar>(i)[j] = 1;
                } else {  // diagonal
                    angles.ptr<uchar>(i)[j] = 2;
                }
            }
        }
    }
}

// step 3
void _non_maximum_suppression(const cv::Mat& dx, const cv::Mat& dy, const cv::Mat& magnitudes, const cv::Mat& angles, cv::Mat& map, std::deque<int>& mapIndicesX, std::deque<int>& mapIndicesY, int low, int high) {
    // 0 -> the pixel may be edge
    // 1 -> the pixel is not edge
    // 2 -> the pixel is edge
    map = cv::Mat::ones(dx.rows + 2, dx.cols + 2, CV_8UC1);
    
    for (int i = 0; i < dx.rows; i++) {
        int r = i + 1;
        uchar* p = map.ptr<uchar>(r);
        for (int j = 0; j < dx.cols; j++) {
            int c = j + 1;
            int m = magnitudes.ptr<int>(r)[c];
            uchar angle = angles.ptr<uchar>(i)[j];
            if (m > low) {
                if (angle == 0) {  // horizontal
                    if (m > magnitudes.ptr<int>(r)[c-1] && m >= magnitudes.ptr<int>(r)[c+1]) {
                        if (m > high) {
                            p[c] = 2;
                            mapIndicesX.push_back(c);
                            mapIndicesY.push_back(r);
                        } else {
                            p[c] = 0;
                        }
                    }
                } else if (angle == 1) {  // vertical
                    if (m > magnitudes.ptr<int>(r-1)[c] && m >= magnitudes.ptr<int>(r+1)[c]) {
                        if (m > high) {
                            p[c] = 2;
                            mapIndicesX.push_back(c);
                            mapIndicesY.push_back(r);
                        } else {
                            p[c] = 0;
                        }
                    }
                } else if (angle == 2) {  // diagonal
                    short xs = dx.ptr<short>(i)[j];
                    short ys = dy.ptr<short>(i)[j];
                    if ((xs > 0 && ys > 0) || (xs < 0 && ys < 0)) {  // 45 degree
                        if (m > magnitudes.ptr<int>(r-1)[c-1] && m > magnitudes.ptr<int>(r+1)[c+1]) {
                            if (m > high) {
                                p[c] = 2;
                                mapIndicesX.push_back(c);
                                mapIndicesY.push_back(r);
                            } else {
                                p[c] = 0;
                            }
                        }
                    } else {  // 135 degree
                        if (m > magnitudes.ptr<int>(r-1)[c+1] && m > magnitudes.ptr<int>(r+1)[c-1]) {
                            if (m > high) {
                                p[c] = 2;
                                mapIndicesX.push_back(c);
                                mapIndicesY.push_back(r);
                            } else {
                                p[c] = 0;
                            }
                        }
                    }
                }
            }
        }
    }
}

// step 4
void _hysteresis_thresholding(std::deque<int>& mapIndicesX, std::deque<int>& mapIndicesY, cv::Mat& map) {
    while (!mapIndicesX.empty()) {
        int r = mapIndicesY.back();
        int c = mapIndicesX.back();
        mapIndicesX.pop_back();
        mapIndicesY.pop_back();
        
        // top left
        if (map.ptr<uchar>(r-1)[c-1] == 0) {
            mapIndicesX.push_back(c-1);
            mapIndicesY.push_back(r-1);
            map.ptr<uchar>(r-1)[c-1] = 2;
        }
        // top
        if (map.ptr<uchar>(r-1)[c] == 0) {
            mapIndicesX.push_back(c);
            mapIndicesY.push_back(r-1);
            map.ptr<uchar>(r-1)[c] = 2;
        }
        // top right
        if (map.ptr<uchar>(r-1)[c+1] == 0) {
            mapIndicesX.push_back(c+1);
            mapIndicesY.push_back(r-1);
            map.ptr<uchar>(r-1)[c+1] = 2;
        }
        // left
        if (map.ptr<uchar>(r)[c-1] == 0) {
            mapIndicesX.push_back(c-1);
            mapIndicesY.push_back(r);
            map.ptr<uchar>(r)[c-1] = 2;
        }
        // right
        if (map.ptr<uchar>(r)[c+1] == 0) {
            mapIndicesX.push_back(c+1);
            mapIndicesY.push_back(r);
            map.ptr<uchar>(r)[c+1] = 2;
        }
        // bottom left
        if (map.ptr<uchar>(r+1)[c-1] == 0) {
            mapIndicesX.push_back(c-1);
            mapIndicesY.push_back(r+1);
            map.ptr<uchar>(r+1)[c-1] = 2;
        }
        // bottom
        if (map.ptr<uchar>(r+1)[c] == 0) {
            mapIndicesX.push_back(c);
            mapIndicesY.push_back(r+1);
            map.ptr<uchar>(r+1)[c] = 2;
        }
        // bottom right
        if (map.ptr<uchar>(r+1)[c+1] == 0) {
            mapIndicesX.push_back(c+1);
            mapIndicesY.push_back(r+1);
            map.ptr<uchar>(r+1)[c+1] = 2;
        }
    }
}

cv::Mat _get_canny_result(const cv::Mat& map) {
    cv::Mat dst(map.rows-2, map.cols-2, CV_8UC1);
    for (int i = 0; i < dst.rows; i++) {
        for (int j = 0; j < dst.cols; j++) {
            dst.ptr<uchar>(i)[j] = (map.ptr<uchar>(i+1)[j+1] == 2 ? 255 : 0);
        }
    }
    return dst;
}

cv::Mat my_canny(const cv::Mat& src, double lowThresh, double highThresh, int apertureSize, bool L2gradient) {
    CV_Assert(src.type() == CV_8UC1);
    CV_Assert(apertureSize == 3 ||  apertureSize == 5);
    
    if (lowThresh > highThresh) {
        std::swap(lowThresh, highThresh);
    }
    
    if (L2gradient) {
        lowThresh = std::min(32767.0, lowThresh);
        highThresh = std::min(32767.0, highThresh);
        if (lowThresh > 0) {
            lowThresh *= lowThresh;
        }
        if (highThresh > 0) {
            highThresh *= highThresh;
        }
    }
    int low = std::floor(lowThresh);
    int high = std::floor(highThresh);
    
    // step 1 : gaussian blur
    cv::Mat gaussian = _gaussian_filter(src);
    
    // step 2: sobel
    cv::Mat dx, dy, magnitudes, angles;
    _sobel_gradient(gaussian, dx, dy, magnitudes, angles, apertureSize, L2gradient);
    
    // step 3: non maximum suppression
    cv::Mat map;
    std::deque<int> mapIndicesX, mapIndicesY;
    _non_maximum_suppression(dx, dy, magnitudes, angles, map, mapIndicesX, mapIndicesY, low, high);
    
    // step 4: hysteresis thresholding
    _hysteresis_thresholding(mapIndicesX, mapIndicesY, map);
    
    cv::Mat dst = _get_canny_result(map);
    
    return dst;
}
        

int main(int argc, char **argv) {
    cv::Mat mat = cv::imread("/home/xlll/Downloads/opencv/samples/data/lena.jpg", cv::IMREAD_GRAYSCALE);
    if (mat.empty()) {
        std::cout << "failed to read image" << std::endl;
        return EXIT_FAILURE;
    }
    
    cv::Mat canny;
    double ratio = 3;
    double threshold1 = 80, threshold2 = threshold1 * ratio;
    int apertureSize = 3;
    bool L2gradient = false;
    cv::Mat gaussian = _gaussian_filter(mat);
    cv::Canny(gaussian, canny, threshold1, threshold2, apertureSize, L2gradient);
    cv::Mat canny2 = my_canny(mat, threshold1, threshold2, apertureSize, L2gradient);
    
    cv::Mat diff = canny != canny2;
    std::vector<cv::Point> nonzeros;
    cv::findNonZero(diff, nonzeros);
    std::cout << "nonzeros size = " << nonzeros.size() << std::endl;
    
    cv::imshow("mat", mat);
    cv::imshow("canny", canny);
    cv::imshow("canny2", canny2);
    cv::waitKey(0);
    
    
    return EXIT_SUCCESS;
}
